package cbr

import (
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"log"
	"sync"
	"time"
)

// CBRPhase represents the phases of CBR protocol
type CBRPhase int

const (
	PhaseInit CBRPhase = iota
	PhasePrepare
	PhaseShareReduce
	PhaseProactivize
	PhaseShareDist
	PhaseComplete
)

func (p CBRPhase) String() string {
	switch p {
	case PhaseInit:
		return "Init"
	case PhasePrepare:
		return "Prepare"
	case PhaseShareReduce:
		return "ShareReduce"
	case PhaseProactivize:
		return "Proactivize"
	case PhaseShareDist:
		return "ShareDist"
	case PhaseComplete:
		return "Complete"
	default:
		return "Unknown"
	}
}

// CBRNode represents a node participating in CBR
type CBRNode struct {
	ID              int
	TotalNodes      int
	FaultyNodes     int
	CurrentPhase    CBRPhase
	Epoch           int
	Membership      []string
	SecretShares    map[int][]byte
	ReceivedShares  map[int][]byte
	ProactivizedKey []byte
	logger          *log.Logger
	mu              sync.RWMutex
}

// NewCBRNode creates a new CBR node
func NewCBRNode(id, totalNodes, faultyNodes int, logger *log.Logger) *CBRNode {
	if logger == nil {
		logger = log.Default()
	}
	
	return &CBRNode{
		ID:             id,
		TotalNodes:     totalNodes,
		FaultyNodes:    faultyNodes,
		CurrentPhase:   PhaseInit,
		SecretShares:   make(map[int][]byte),
		ReceivedShares: make(map[int][]byte),
		logger:         logger,
	}
}

// Execute runs the full CBR protocol
func (n *CBRNode) Execute(epoch int, newMembership []string) error {
	n.mu.Lock()
	n.Epoch = epoch
	n.Membership = newMembership
	n.mu.Unlock()
	
	n.logger.Printf("[CBR Node %d] Starting CBR protocol for epoch %d", n.ID, epoch)
	
	// Phase 1: Prepare
	if err := n.executePrepare(); err != nil {
		return fmt.Errorf("prepare phase failed: %w", err)
	}
	
	// Phase 2: ShareReduce
	if err := n.executeShareReduce(); err != nil {
		return fmt.Errorf("share reduce phase failed: %w", err)
	}
	
	// Phase 3: Proactivize
	if err := n.executeProactivize(); err != nil {
		return fmt.Errorf("proactivize phase failed: %w", err)
	}
	
	// Phase 4: ShareDist
	if err := n.executeShareDist(); err != nil {
		return fmt.Errorf("share dist phase failed: %w", err)
	}
	
	n.mu.Lock()
	n.CurrentPhase = PhaseComplete
	n.mu.Unlock()
	
	n.logger.Printf("[CBR Node %d] CBR protocol completed for epoch %d", n.ID, epoch)
	return nil
}

// executePrepare executes the Prepare phase
func (n *CBRNode) executePrepare() error {
	n.mu.Lock()
	n.CurrentPhase = PhasePrepare
	n.mu.Unlock()
	
	n.logger.Printf("[CBR Node %d] Phase 1: Prepare", n.ID)
	
	// Generate initial secret shares for all nodes
	for i := 0; i < n.TotalNodes; i++ {
		share := n.generateSecretShare(i)
		n.SecretShares[i] = share
	}
	
	// Simulate network delay
	time.Sleep(10 * time.Millisecond)
	
	return nil
}

// executeShareReduce executes the ShareReduce phase
func (n *CBRNode) executeShareReduce() error {
	n.mu.Lock()
	n.CurrentPhase = PhaseShareReduce
	n.mu.Unlock()
	
	n.logger.Printf("[CBR Node %d] Phase 2: ShareReduce", n.ID)
	
	// Simulate receiving shares from other nodes
	for i := 0; i < n.TotalNodes; i++ {
		if i != n.ID {
			// In real implementation, this would be received over network
			share := n.generateSecretShare(i)
			n.ReceivedShares[i] = share
		}
	}
	
	// Reduce shares to create new configuration
	reducedShare := n.reduceShares()
	n.SecretShares[n.ID] = reducedShare
	
	// Simulate network delay
	time.Sleep(15 * time.Millisecond)
	
	return nil
}

// executeProactivize executes the Proactivize phase
func (n *CBRNode) executeProactivize() error {
	n.mu.Lock()
	n.CurrentPhase = PhaseProactivize
	n.mu.Unlock()
	
	n.logger.Printf("[CBR Node %d] Phase 3: Proactivize", n.ID)
	
	// Proactivize the secret key with fresh randomness
	n.ProactivizedKey = n.proactivizeKey()
	
	// Simulate computation delay
	time.Sleep(20 * time.Millisecond)
	
	return nil
}

// executeShareDist executes the ShareDist phase
func (n *CBRNode) executeShareDist() error {
	n.mu.Lock()
	n.CurrentPhase = PhaseShareDist
	n.mu.Unlock()
	
	n.logger.Printf("[CBR Node %d] Phase 4: ShareDist", n.ID)
	
	// Distribute new shares to all nodes in new membership
	for _, member := range n.Membership {
		// In real implementation, send share to member over network
		_ = member
	}
	
	// Simulate network delay
	time.Sleep(10 * time.Millisecond)
	
	return nil
}

// Helper functions

func (n *CBRNode) generateSecretShare(nodeID int) []byte {
	// Generate deterministic secret share based on node ID and epoch
	data := fmt.Sprintf("node-%d-epoch-%d-share-%d", n.ID, n.Epoch, nodeID)
	hash := sha256.Sum256([]byte(data))
	return hash[:]
}

func (n *CBRNode) reduceShares() []byte {
	// Combine all received shares using XOR (simplified)
	combined := make([]byte, 32)
	
	for _, share := range n.ReceivedShares {
		for i := 0; i < len(combined) && i < len(share); i++ {
			combined[i] ^= share[i]
		}
	}
	
	// Mix with own share
	if ownShare, ok := n.SecretShares[n.ID]; ok {
		for i := 0; i < len(combined) && i < len(ownShare); i++ {
			combined[i] ^= ownShare[i]
		}
	}
	
	return combined
}

func (n *CBRNode) proactivizeKey() []byte {
	// Add fresh randomness to the key (simplified)
	freshRandomness := fmt.Sprintf("fresh-random-node-%d-epoch-%d-%d", 
		n.ID, n.Epoch, time.Now().UnixNano())
	hash := sha256.Sum256([]byte(freshRandomness))
	
	// Combine with existing share
	proactivized := make([]byte, 32)
	ownShare := n.SecretShares[n.ID]
	
	for i := 0; i < len(proactivized); i++ {
		if i < len(ownShare) {
			proactivized[i] = ownShare[i] ^ hash[i]
		} else {
			proactivized[i] = hash[i]
		}
	}
	
	return proactivized
}

// GetCurrentPhase returns the current phase
func (n *CBRNode) GetCurrentPhase() CBRPhase {
	n.mu.RLock()
	defer n.mu.RUnlock()
	return n.CurrentPhase
}

// GetProactivizedKey returns the proactivized key
func (n *CBRNode) GetProactivizedKey() []byte {
	n.mu.RLock()
	defer n.mu.RUnlock()
	return n.ProactivizedKey
}

// MarshalJSON implements json.Marshaler
func (n *CBRNode) MarshalJSON() ([]byte, error) {
	n.mu.RLock()
	defer n.mu.RUnlock()
	
	return json.Marshal(map[string]interface{}{
		"id":            n.ID,
		"total_nodes":   n.TotalNodes,
		"faulty_nodes":  n.FaultyNodes,
		"current_phase": n.CurrentPhase.String(),
		"epoch":         n.Epoch,
		"membership":    n.Membership,
	})
}

